package com.fortune.cat.filter.factory;

import com.fortune.cat.bean.RequestRateLimiterConfig;
import com.fortune.cat.enums.RateLimiterEnums;
import com.fortune.cat.limiter.LocalCacheRateLimiter;
import com.fortune.cat.resolver.AbstractKeyResolver;
import com.fortune.cat.utils.ServerResponseUtils;
import com.google.gson.Gson;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.ratelimit.RateLimiter;
import org.springframework.cloud.gateway.filter.ratelimit.RedisRateLimiter;
import org.springframework.cloud.gateway.route.Route;
import org.springframework.cloud.gateway.support.ServerWebExchangeUtils;
import org.springframework.http.HttpStatus;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

import java.util.HashMap;
import java.util.Map;

import static org.springframework.cloud.gateway.support.ServerWebExchangeUtils.GATEWAY_PREDICATE_ROUTE_ATTR;

/**
 * @author zhaozeyang
 * @date 2019/4/24 21:29
 * @description 默认的限流拦截器
 */
public class DefaultRequestRateLimiterFilterFactory extends AbstractCustomizeGatewayFilterFactory<RequestRateLimiterConfig> {


    private final AbstractKeyResolver keyResolver;

    private final RedisRateLimiter redisRateLimiter;

    private final LocalCacheRateLimiter localCacheRateLimiter;

    public static final String DEFAULT_RATE_LIMITER_FILTER_NAME = "DefaultRequestRateLimiterFilterFactory";

    public static final String EMPTY_KEY = "____EMPTY_KEY__";

    public static final Gson GSON = new Gson();

    private static String ERROR_RETURN_STR;

    static {
        Map<String, Object> map = new HashMap<>();
        map.put("status", 403000);
        map.put("msg", "TOO_MANY_REQUESTS");
        ERROR_RETURN_STR = GSON.toJson(map);
    }

    public DefaultRequestRateLimiterFilterFactory(AbstractKeyResolver keyResolver, RedisRateLimiter redisRateLimiter, LocalCacheRateLimiter localCacheRateLimiter) {
        super(RequestRateLimiterConfig.class);
        this.keyResolver = keyResolver;
        this.redisRateLimiter = redisRateLimiter;
        this.localCacheRateLimiter = localCacheRateLimiter;
    }

    @Override
    protected Mono<Void> doApply(RequestRateLimiterConfig config, ServerWebExchange exchange, GatewayFilterChain chain) {
        keyResolver.setRequestRateLimiterConfig(config);
        String routeId = ((Route)exchange.getAttribute(ServerWebExchangeUtils.GATEWAY_ROUTE_ATTR)).getId();
        return keyResolver.resolve(exchange).defaultIfEmpty(EMPTY_KEY).flatMap(key -> {
            if (EMPTY_KEY.equals(key)) {
                return chain.filter(exchange);
            }
            Integer type = config.getRateLimiterType();
            if (RateLimiterEnums.LOCAL_RATE_LIMITER.getType().equals(type)) {
                localCacheRateLimiter.setRequestRateLimiterConfig(config);
                return responseMono(localCacheRateLimiter.isAllowed(key), exchange, chain);
            } else if (RateLimiterEnums.REDIS_RATE_LIMITER.getType().equals(type)) {
                return responseMono(redisRateLimiter.isAllowed(routeId, key), exchange, chain);
            } else {
                return chain.filter(exchange);
            }
        });
    }

    private Mono<Void> responseMono(Mono<RateLimiter.Response> responseMono, ServerWebExchange exchange, GatewayFilterChain chain) {
        return responseMono.flatMap(response -> {
            for (Map.Entry<String, String> header : response.getHeaders()
                    .entrySet()) {
                exchange.getResponse().getHeaders().add(header.getKey(),
                        header.getValue());
            }

            if (response.isAllowed()) {
                return chain.filter(exchange);
            }
            return ServerResponseUtils.returnCommonResponse(ERROR_RETURN_STR, exchange, HttpStatus.TOO_MANY_REQUESTS);
        });
    }


    @Override
    public String name() {
        return DEFAULT_RATE_LIMITER_FILTER_NAME;
    }

    @Override
    protected int doGetOrder() {

        return 0;
    }

    @Override
    protected boolean skipPresent(RequestRateLimiterConfig requestRateLimiterConfig, ServerWebExchange exchange) {
        return false;
    }

}
